import os
import random
from dotenv import load_dotenv
from volcenginesdkarkruntime import Ark

# 加载环境变量
load_dotenv()

# 初始化Ark客户端
client = Ark(
    base_url="https://ark.cn-beijing.volces.com/api/v3",
    api_key=os.environ.get("ARK_API_KEY"),
)


def generate_image(prompt: str, image_url: str) -> str:
    """
    根据提示词和原图生成新图片

    Args:
        prompt (str): 修改提示词
        image_url (str): 原始图片的URL

    Returns:
        str: 生成图片的URL

    Raises:
        Exception: 图片生成失败时抛出异常
    """
    try:
        # 调用API生成图片，使用固定的参数配置
        images_response = client.images.generate(
            model="doubao-seededit-3-0-i2i-250628",
            prompt=prompt,
            image=image_url,
            seed=random.randint(1, 10000),
            guidance_scale=8.5,
            size="adaptive",
            watermark=False
        )

        # 检查响应并返回图片URL
        if images_response.data and len(images_response.data) > 0:
            return images_response.data[0].url
        else:
            raise Exception("图片生成失败：未返回有效的图片数据")

    except Exception as e:
        raise Exception(f"图片生成过程中发生错误: {str(e)}")


# 使用示例
if __name__ == "__main__":
    try:
        result_url = generate_image(
            prompt="改成爱心形状的泡泡",
            image_url="https://ark-project.tos-cn-beijing.volces.com/doc_image/seededit_i2i.jpeg"
        )
        print(f"生成的图片URL: {result_url}")
    except Exception as e:
        print(f"错误: {e}")
